"""
glass 反照率数据线性插值
"""
import os
import time

import numpy as np
from osgeo import gdal, osr
from tqdm import tqdm


# 读取tif文件，输出图像数组和仿射变换矩阵
def read_tif(file_path):
    dataset = gdal.Open(file_path)
    # 读成数组
    data = dataset.ReadAsArray()
    # 得到仿射变换矩阵
    geotrans = dataset.GetGeoTransform()
    return data, geotrans


# 写入tif
def write_tif(file_path, data, geotransform, nodata, dataType):
    gdal_type = ''
    if dataType == 'int16':
        gdal_type = gdal.GDT_Int16
    elif dataType == 'int32':
        gdal_type = gdal.GDT_Int32
    elif dataType == 'float32':
        gdal_type = gdal.GDT_Float32
    elif dataType == 'float64':
        gdal_type = gdal.GDT_Float64
    elif dataType == 'byte':
        gdal_type = gdal.GDT_Byte
    elif dataType == 'uint16':
        gdal_type = gdal.GDT_UInt16
    elif dataType == 'uint32':
        gdal_type = gdal.GDT_UInt32
    # 相当于一个创建数据集的驱动
    driver = gdal.GetDriverByName("GTiff")
    # 根据数据的维度来确定行列数（图像大小）
    rows, cols = data.shape
    # 创建一个新的数据集，存储输出文件；1是波段
    dataset = driver.Create(file_path, cols, rows, 1, gdal_type)
    # 设置仿射变换矩阵
    dataset.SetGeoTransform(geotransform)
    # 设置投影
    prj = osr.SpatialReference()
    # WGS84投影
    prj.ImportFromEPSG(4326)
    dataset.SetProjection(prj.ExportToWkt())
    # 1个波段
    band = dataset.GetRasterBand(1)
    band.WriteArray(data)
    band.SetNoDataValue(nodata)
    # 释放内存
    del dataset


def glass_linear(glass_dir, output_dir):
    start = time.time()
    basename = ".tif"
    glass_list = [os.path.join(glass_dir, f) for f in os.listdir(glass_dir) if f.endswith(basename)]
    date_list = [os.path.splitext(os.path.basename(f))[0] for f in glass_list]
    # print(date_list)
    time_array = np.arange(int(date_list[0][-3:]), int(date_list[-1][-3:]) + 1, 8)
    # print(time_array)
    dataset = gdal.Open(glass_list[0])
    rows, cols = dataset.RasterYSize, dataset.RasterXSize
    geotransform = dataset.GetGeoTransform()
    nodata = dataset.GetRasterBand(1).GetNoDataValue()
    glass_array = np.zeros((len(glass_list), rows, cols), dtype=np.int16)
    for i in range(len(glass_list)):
        data, _ = read_tif(glass_list[i])
        glass_array[i, :, :] = data

    inter_date = []
    for date in date_list:
        while str(int(date) + 1) not in date_list and int(date) + 1 <= int(date_list[-1]):
            inter_date.append(str(int(date) + 1))
            date = str(int(date) + 1)
    print(inter_date)
    inter_time = [int(date[-3:]) for date in inter_date]

    result_array = np.zeros((rows, cols, len(inter_date)), dtype=np.int16)

    glass_array = np.transpose(glass_array, (1, 2, 0))

    for i in tqdm(range(glass_array.shape[0])):
        for j in range(glass_array.shape[1]):
            time_series = glass_array[i, j, :]
            # 线性内插
            result_array[i, j, :] = np.interp(inter_time, time_array, time_series)

    result_array = np.transpose(result_array, (2, 0, 1))
    result_array.dtype = np.int16
    for i in range(result_array.shape[0]):
        output_image = output_dir + "/" + inter_date[i] + ".tif"
        write_tif(output_image, result_array[i, :, :], geotransform, nodata=nodata, dataType='int16')
    end = time.time()
    print("完成插值，耗时：{}s".format(end - start))


if __name__ == '__main__':
    glass_dir = r"G:\test\process_result\GLASS_tif\merge_clip\QC_VIS"
    output_dir = r"G:\test\process_result\GLASS_tif\glass_linear\QC_VIS"
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    glass_linear(glass_dir, output_dir)
